{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "accelerate config\n",
    "accelerate launch train_text_to_image_lora.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "2025-04-05 15:35:10.319444: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.\n",
      "2025-04-05 15:35:10.371688: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2025-04-05 15:35:10.371721: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2025-04-05 15:35:10.371751: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2025-04-05 15:35:10.382075: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.\n",
      "2025-04-05 15:35:10.383101: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2025-04-05 15:35:11.691624: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "import logging\n",
    "import math\n",
    "import os\n",
    "import random\n",
    "import shutil\n",
    "from contextlib import nullcontext\n",
    "from pathlib import Path\n",
    "\n",
    "import datasets\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.checkpoint\n",
    "import transformers\n",
    "from accelerate import Accelerator\n",
    "from accelerate.logging import get_logger\n",
    "from accelerate.utils import ProjectConfiguration, set_seed\n",
    "from datasets import load_dataset\n",
    "from huggingface_hub import create_repo, upload_folder\n",
    "from packaging import version\n",
    "from peft import LoraConfig\n",
    "from peft.utils import get_peft_model_state_dict\n",
    "from torchvision import transforms\n",
    "from tqdm.auto import tqdm\n",
    "from transformers import CLIPTextModel, CLIPTokenizer\n",
    "\n",
    "import diffusers\n",
    "from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, StableDiffusionPipeline, UNet2DConditionModel\n",
    "from diffusers.optimization import get_scheduler\n",
    "from diffusers.training_utils import cast_training_params, compute_snr\n",
    "from diffusers.utils import check_min_version, convert_state_dict_to_diffusers, is_wandb_available\n",
    "from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card\n",
    "from diffusers.utils.import_utils import is_xformers_available\n",
    "from diffusers.utils.torch_utils import is_compiled_module\n",
    "\n",
    "\n",
    "if is_wandb_available():\n",
    "    import wandb\n",
    "\n",
    "# Will error if the minimal version of diffusers is not installed. Remove at your own risks.\n",
    "#check_min_version(\"0.33.0.dev0\")\n",
    "\n",
    "logger = get_logger(__name__, log_level=\"INFO\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def save_model_card(\n",
    "    repo_id: str,\n",
    "    images: list = None,\n",
    "    base_model: str = None,\n",
    "    dataset_name: str = None,\n",
    "    repo_folder: str = None,\n",
    "):\n",
    "    img_str = \"\"\n",
    "    if images is not None:\n",
    "        for i, image in enumerate(images):\n",
    "            image.save(os.path.join(repo_folder, f\"image_{i}.png\"))\n",
    "            img_str += f\"![img_{i}](./image_{i}.png)\\n\"\n",
    "\n",
    "    model_description = f\"\"\"\n",
    "# LoRA text2image fine-tuning - {repo_id}\n",
    "These are LoRA adaption weights for {base_model}. The weights were fine-tuned on the {dataset_name} dataset. You can find some example images in the following. \\n\n",
    "{img_str}\n",
    "\"\"\"\n",
    "\n",
    "    model_card = load_or_create_model_card(\n",
    "        repo_id_or_path=repo_id,\n",
    "        from_training=True,\n",
    "        license=\"creativeml-openrail-m\",\n",
    "        base_model=base_model,\n",
    "        model_description=model_description,\n",
    "        inference=True,\n",
    "    )\n",
    "\n",
    "    tags = [\n",
    "        \"stable-diffusion\",\n",
    "        \"stable-diffusion-diffusers\",\n",
    "        \"text-to-image\",\n",
    "        \"diffusers\",\n",
    "        \"diffusers-training\",\n",
    "        \"lora\",\n",
    "    ]\n",
    "    model_card = populate_model_card(model_card, tags=tags)\n",
    "\n",
    "    model_card.save(os.path.join(repo_folder, \"README.md\"))\n",
    "\n",
    "\n",
    "def log_validation(\n",
    "    pipeline,\n",
    "    args,\n",
    "    accelerator,\n",
    "    epoch,\n",
    "    is_final_validation=False,\n",
    "):\n",
    "    logger.info(\n",
    "        f\"Running validation... \\n Generating {args.num_validation_images} images with prompt:\"\n",
    "        f\" {args.validation_prompt}.\"\n",
    "    )\n",
    "    pipeline = pipeline.to(accelerator.device)\n",
    "    pipeline.set_progress_bar_config(disable=True)\n",
    "    generator = torch.Generator(device=accelerator.device)\n",
    "    if args.seed is not None:\n",
    "        generator = generator.manual_seed(args.seed)\n",
    "    images = []\n",
    "    if torch.backends.mps.is_available():\n",
    "        autocast_ctx = nullcontext()\n",
    "    else:\n",
    "        autocast_ctx = torch.autocast(accelerator.device.type)\n",
    "\n",
    "    with autocast_ctx:\n",
    "        for _ in range(args.num_validation_images):\n",
    "            images.append(pipeline(args.validation_prompt, num_inference_steps=30, generator=generator).images[0])\n",
    "\n",
    "    for tracker in accelerator.trackers:\n",
    "        phase_name = \"test\" if is_final_validation else \"validation\"\n",
    "        if tracker.name == \"tensorboard\":\n",
    "            np_images = np.stack([np.asarray(img) for img in images])\n",
    "            tracker.writer.add_images(phase_name, np_images, epoch, dataformats=\"NHWC\")\n",
    "        if tracker.name == \"wandb\":\n",
    "            tracker.log(\n",
    "                {\n",
    "                    phase_name: [\n",
    "                        wandb.Image(image, caption=f\"{i}: {args.validation_prompt}\") for i, image in enumerate(images)\n",
    "                    ]\n",
    "                }\n",
    "            )\n",
    "    return images\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass, field\n",
    "from typing import Optional\n",
    "\n",
    "@dataclass\n",
    "class TrainingArguments:\n",
    "    pretrained_model_name_or_path: str = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"Path to pretrained model or model identifier from huggingface.co/models.\"}\n",
    "    )\n",
    "    revision: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"Revision of pretrained model identifier from huggingface.co/models.\"}\n",
    "    )\n",
    "    variant: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16\"}\n",
    "    )\n",
    "    dataset_name: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"The name of the Dataset (from the HuggingFace hub) to train on.\"}\n",
    "    )\n",
    "    dataset_config_name: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"The config of the Dataset, leave as None if there's only one config.\"}\n",
    "    )\n",
    "    train_data_dir: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"A folder containing the training data.\"}\n",
    "    )\n",
    "    image_column: str = field(\n",
    "        default=\"image\",\n",
    "        metadata={\"help\": \"The column of the dataset containing an image.\"}\n",
    "    )\n",
    "    caption_column: str = field(\n",
    "        default=\"text\",\n",
    "        metadata={\"help\": \"The column of the dataset containing a caption or a list of captions.\"}\n",
    "    )\n",
    "    validation_prompt: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"A prompt that is sampled during training for inference.\"}\n",
    "    )\n",
    "    num_validation_images: int = field(\n",
    "        default=4,\n",
    "        metadata={\"help\": \"Number of images that should be generated during validation with `validation_prompt`.\"}\n",
    "    )\n",
    "    validation_epochs: int = field(\n",
    "        default=1,\n",
    "        metadata={\"help\": \"Run fine-tuning validation every X epochs.\"}\n",
    "    )\n",
    "    max_train_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"For debugging purposes or quicker training, truncate the number of training examples.\"}\n",
    "    )\n",
    "    output_dir: str = field(\n",
    "        default=\"sd-model-finetuned-lora\",\n",
    "        metadata={\"help\": \"The output directory where the model predictions and checkpoints will be written.\"}\n",
    "    )\n",
    "    cache_dir: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"The directory where the downloaded models and datasets will be stored.\"}\n",
    "    )\n",
    "    seed: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"A seed for reproducible training.\"}\n",
    "    )\n",
    "    resolution: int = field(\n",
    "        default=512,\n",
    "        metadata={\"help\": \"The resolution for input images.\"}\n",
    "    )\n",
    "    center_crop: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether to center crop the input images to the resolution.\"}\n",
    "    )\n",
    "    random_flip: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether to randomly flip images horizontally.\"}\n",
    "    )\n",
    "    train_batch_size: int = field(\n",
    "        default=16,\n",
    "        metadata={\"help\": \"Batch size (per device) for the training dataloader.\"}\n",
    "    )\n",
    "    num_train_epochs: int = field(\n",
    "        default=100,\n",
    "        metadata={\"help\": \"Total number of training epochs.\"}\n",
    "    )\n",
    "    max_train_steps: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"Total number of training steps to perform. Overrides num_train_epochs if provided.\"}\n",
    "    )\n",
    "    gradient_accumulation_steps: int = field(\n",
    "        default=1,\n",
    "        metadata={\"help\": \"Number of updates steps to accumulate before performing a backward/update pass.\"}\n",
    "    )\n",
    "    gradient_checkpointing: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether or not to use gradient checkpointing to save memory.\"}\n",
    "    )\n",
    "    learning_rate: float = field(\n",
    "        default=1e-4,\n",
    "        metadata={\"help\": \"Initial learning rate (after the potential warmup period).\"}\n",
    "    )\n",
    "    scale_lr: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.\"}\n",
    "    )\n",
    "    lr_scheduler: str = field(\n",
    "        default=\"constant\",\n",
    "        metadata={\"help\": \"The scheduler type to use.\"}\n",
    "    )\n",
    "    lr_warmup_steps: int = field(\n",
    "        default=500,\n",
    "        metadata={\"help\": \"Number of steps for the warmup in the lr scheduler.\"}\n",
    "    )\n",
    "    snr_gamma: Optional[float] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"SNR weighting gamma to be used if rebalancing the loss.\"}\n",
    "    )\n",
    "    use_8bit_adam: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether or not to use 8-bit Adam from bitsandbytes.\"}\n",
    "    )\n",
    "    allow_tf32: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether or not to allow TF32 on Ampere GPUs.\"}\n",
    "    )\n",
    "    dataloader_num_workers: int = field(\n",
    "        default=0,\n",
    "        metadata={\"help\": \"Number of subprocesses to use for data loading.\"}\n",
    "    )\n",
    "    adam_beta1: float = field(\n",
    "        default=0.9,\n",
    "        metadata={\"help\": \"The beta1 parameter for the Adam optimizer.\"}\n",
    "    )\n",
    "    adam_beta2: float = field(\n",
    "        default=0.999,\n",
    "        metadata={\"help\": \"The beta2 parameter for the Adam optimizer.\"}\n",
    "    )\n",
    "    adam_weight_decay: float = field(\n",
    "        default=1e-2,\n",
    "        metadata={\"help\": \"Weight decay to use.\"}\n",
    "    )\n",
    "    adam_epsilon: float = field(\n",
    "        default=1e-08,\n",
    "        metadata={\"help\": \"Epsilon value for the Adam optimizer.\"}\n",
    "    )\n",
    "    max_grad_norm: float = field(\n",
    "        default=1.0,\n",
    "        metadata={\"help\": \"Max gradient norm.\"}\n",
    "    )\n",
    "    push_to_hub: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether or not to push the model to the Hub.\"}\n",
    "    )\n",
    "    hub_token: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"The token to use to push to the Model Hub.\"}\n",
    "    )\n",
    "    prediction_type: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"The prediction_type that shall be used for training.\"}\n",
    "    )\n",
    "    hub_model_id: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"The name of the repository to keep in sync with the local `output_dir`.\"}\n",
    "    )\n",
    "    logging_dir: str = field(\n",
    "        default=\"logs\",\n",
    "        metadata={\"help\": \"TensorBoard log directory.\"}\n",
    "    )\n",
    "    mixed_precision: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"Whether to use mixed precision. Choose between fp16 and bf16.\"}\n",
    "    )\n",
    "    report_to: str = field(\n",
    "        default=\"tensorboard\",\n",
    "        metadata={\"help\": \"The integration to report the results and logs to.\"}\n",
    "    )\n",
    "    local_rank: int = field(\n",
    "        default=-1,\n",
    "        metadata={\"help\": \"For distributed training: local_rank.\"}\n",
    "    )\n",
    "    checkpointing_steps: int = field(\n",
    "        default=500,\n",
    "        metadata={\"help\": \"Save a checkpoint of the training state every X updates.\"}\n",
    "    )\n",
    "    checkpoints_total_limit: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"Max number of checkpoints to store.\"}\n",
    "    )\n",
    "    resume_from_checkpoint: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"Whether training should be resumed from a previous checkpoint.\"}\n",
    "    )\n",
    "    enable_xformers_memory_efficient_attention: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether or not to use xformers.\"}\n",
    "    )\n",
    "    noise_offset: float = field(\n",
    "        default=0,\n",
    "        metadata={\"help\": \"The scale of noise offset.\"}\n",
    "    )\n",
    "    rank: int = field(\n",
    "        default=4,\n",
    "        metadata={\"help\": \"The dimension of the LoRA update matrices.\"}\n",
    "    )\n",
    "\n",
    "\n",
    "import json\n",
    "with open(\"train_sd1.5_lora_config.json\", \"r\") as f:\n",
    "    config = json.load(f)\n",
    "\n",
    "# 实例化 TrainingArguments\n",
    "args = TrainingArguments(**config)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "04/05/2025 15:36:42 - INFO - __main__ - Distributed environment: NO\n",
      "Num processes: 1\n",
      "Process index: 0\n",
      "Local process index: 0\n",
      "Device: cuda\n",
      "\n",
      "Mixed precision type: no\n",
      "\n"
     ]
    }
   ],
   "source": [
    "DATASET_NAME_MAPPING = {\n",
    "    \"lambdalabs/naruto-blip-captions\": (\"image\", \"text\"),\n",
    "}\n",
    "if args.report_to == \"wandb\" and args.hub_token is not None:\n",
    "    raise ValueError(\n",
    "        \"You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token.\"\n",
    "        \" Please use `huggingface-cli login` to authenticate with the Hub.\"\n",
    "    )\n",
    "\n",
    "logging_dir = Path(args.output_dir, args.logging_dir)\n",
    "\n",
    "accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)\n",
    "\n",
    "accelerator = Accelerator(\n",
    "    gradient_accumulation_steps=args.gradient_accumulation_steps,\n",
    "    mixed_precision=args.mixed_precision,\n",
    "    log_with=args.report_to,\n",
    "    project_config=accelerator_project_config,\n",
    ")\n",
    "\n",
    "# Disable AMP for MPS.\n",
    "if torch.backends.mps.is_available():\n",
    "    accelerator.native_amp = False\n",
    "\n",
    "# Make one log on every process with the configuration for debugging.\n",
    "logging.basicConfig(\n",
    "    format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",\n",
    "    datefmt=\"%m/%d/%Y %H:%M:%S\",\n",
    "    level=logging.INFO,\n",
    ")\n",
    "logger.info(accelerator.state, main_process_only=False)\n",
    "if accelerator.is_local_main_process:\n",
    "    datasets.utils.logging.set_verbosity_warning()\n",
    "    transformers.utils.logging.set_verbosity_warning()\n",
    "    diffusers.utils.logging.set_verbosity_info()\n",
    "else:\n",
    "    datasets.utils.logging.set_verbosity_error()\n",
    "    transformers.utils.logging.set_verbosity_error()\n",
    "    diffusers.utils.logging.set_verbosity_error()\n",
    "\n",
    "# If passed along, set the training seed now.\n",
    "if args.seed is not None:\n",
    "    set_seed(args.seed)\n",
    "\n",
    "# Handle the repository creation\n",
    "if accelerator.is_main_process:\n",
    "    if args.output_dir is not None:\n",
    "        os.makedirs(args.output_dir, exist_ok=True)\n",
    "\n",
    "    if args.push_to_hub:\n",
    "        repo_id = create_repo(\n",
    "            repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token\n",
    "        ).repo_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "{'rescale_betas_zero_snr', 'dynamic_thresholding_ratio', 'thresholding', 'variance_type', 'prediction_type', 'clip_sample_range', 'sample_max_value', 'timestep_spacing'} was not found in config. Values will be initialized to default values.\n",
      "{'scaling_factor', 'force_upcast'} was not found in config. Values will be initialized to default values.\n",
      "{'addition_embed_type_num_heads', 'time_cond_proj_dim', 'cross_attention_norm', 'time_embedding_dim', 'upcast_attention', 'timestep_post_act', 'only_cross_attention', 'class_embeddings_concat', 'class_embed_type', 'dual_cross_attention', 'resnet_time_scale_shift', 'use_linear_projection', 'reverse_transformer_layers_per_block', 'attention_type', 'encoder_hid_dim', 'mid_block_only_cross_attention', 'mid_block_type', 'addition_embed_type', 'resnet_skip_time_act', 'time_embedding_act_fn', 'num_attention_heads', 'dropout', 'addition_time_embed_dim', 'num_class_embeds', 'conv_out_kernel', 'time_embedding_type', 'encoder_hid_dim_type', 'resnet_out_scale_factor', 'conv_in_kernel', 'projection_class_embeddings_input_dim', 'transformer_layers_per_block'} was not found in config. Values will be initialized to default values.\n"
     ]
    }
   ],
   "source": [
    "# Load scheduler, tokenizer and models.\n",
    "noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder=\"scheduler\")\n",
    "tokenizer = CLIPTokenizer.from_pretrained(\n",
    "    args.pretrained_model_name_or_path, subfolder=\"tokenizer\", revision=args.revision\n",
    ")\n",
    "text_encoder = CLIPTextModel.from_pretrained(\n",
    "    args.pretrained_model_name_or_path, subfolder=\"text_encoder\", revision=args.revision\n",
    ")\n",
    "vae = AutoencoderKL.from_pretrained(\n",
    "    args.pretrained_model_name_or_path, subfolder=\"vae\", revision=args.revision, variant=args.variant\n",
    ")\n",
    "unet = UNet2DConditionModel.from_pretrained(\n",
    "    args.pretrained_model_name_or_path, subfolder=\"unet\", revision=args.revision, variant=args.variant\n",
    ")\n",
    "# freeze parameters of models to save more memory\n",
    "unet.requires_grad_(False)\n",
    "vae.requires_grad_(False)\n",
    "text_encoder.requires_grad_(False)\n",
    "\n",
    "# For mixed precision training we cast all non-trainable weights (vae, non-lora text_encoder and non-lora unet) to half-precision\n",
    "# as these weights are only used for inference, keeping weights in full precision is not required.\n",
    "weight_dtype = torch.float32\n",
    "if accelerator.mixed_precision == \"fp16\":\n",
    "    weight_dtype = torch.float16\n",
    "elif accelerator.mixed_precision == \"bf16\":\n",
    "    weight_dtype = torch.bfloat16\n",
    "\n",
    "unet_lora_config = LoraConfig(\n",
    "    r=args.rank,\n",
    "    lora_alpha=args.rank,\n",
    "    init_lora_weights=\"gaussian\",\n",
    "    target_modules=[\"to_k\", \"to_q\", \"to_v\", \"to_out.0\"],\n",
    ")\n",
    "\n",
    "# Move unet, vae and text_encoder to device and cast to weight_dtype\n",
    "unet.to(accelerator.device, dtype=weight_dtype)\n",
    "vae.to(accelerator.device, dtype=weight_dtype)\n",
    "text_encoder.to(accelerator.device, dtype=weight_dtype)\n",
    "\n",
    "# Add adapter and make sure the trainable params are in float32.\n",
    "unet.add_adapter(unet_lora_config)\n",
    "if args.mixed_precision == \"fp16\":\n",
    "    # only upcast trainable parameters (LoRA) into fp32\n",
    "    cast_training_params(unet, dtype=torch.float32)\n",
    "\n",
    "# xformer 加速attention计算\n",
    "if args.enable_xformers_memory_efficient_attention:\n",
    "    if is_xformers_available():\n",
    "        import xformers\n",
    "\n",
    "        xformers_version = version.parse(xformers.__version__)\n",
    "        if xformers_version == version.parse(\"0.0.16\"):\n",
    "            logger.warning(\n",
    "                \"xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details.\"\n",
    "            )\n",
    "        unet.enable_xformers_memory_efficient_attention()\n",
    "    else:\n",
    "        raise ValueError(\"xformers is not available. Make sure it is installed correctly\")\n",
    "\n",
    "lora_layers = filter(lambda p: p.requires_grad, unet.parameters())\n",
    "if args.gradient_checkpointing:\n",
    "    unet.enable_gradient_checkpointing()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enable TF32 for faster training on Ampere GPUs,eg. A100 H100\n",
    "if args.allow_tf32:\n",
    "    torch.backends.cuda.matmul.allow_tf32 = True\n",
    "\n",
    "if args.scale_lr:\n",
    "    args.learning_rate = (\n",
    "        args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes\n",
    "    )\n",
    "\n",
    "# Initialize the optimizer\n",
    "# 8bit-optimizer will use less GPU memory\n",
    "if args.use_8bit_adam:\n",
    "    try:\n",
    "        import bitsandbytes as bnb\n",
    "    except ImportError:\n",
    "        raise ImportError(\n",
    "            \"Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`\"\n",
    "        )\n",
    "\n",
    "    optimizer_cls = bnb.optim.AdamW8bit\n",
    "else:\n",
    "    optimizer_cls = torch.optim.AdamW\n",
    "\n",
    "optimizer = optimizer_cls(\n",
    "    lora_layers,\n",
    "    lr=args.learning_rate,\n",
    "    betas=(args.adam_beta1, args.adam_beta2),\n",
    "    weight_decay=args.adam_weight_decay,\n",
    "    eps=args.adam_epsilon,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "04/05/2025 15:41:36 - INFO - __main__ - use local dataset\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=716x703 at 0x7F0AFD06BDF0>, 'text': 'Pixel art, 1girl, solo, long hair, skirt, simple background, anime, full body, from side, running, narrowed eyes, long hair, purple hair, hime_cut, skirt, dress, Long Boots'}\n"
     ]
    }
   ],
   "source": [
    "# datasets库读取数据，格式类似于Pandas读取表格\n",
    "if args.dataset_name is not None:\n",
    "    # Downloading and loading a dataset from the hub.\n",
    "    logger.info(\"download dataset from hub\")\n",
    "    dataset = load_dataset(\n",
    "        args.dataset_name,\n",
    "        args.dataset_config_name,\n",
    "        cache_dir=args.cache_dir,\n",
    "        data_dir=args.train_data_dir,\n",
    "    )\n",
    "else:\n",
    "    # loading a local dataset.\n",
    "    logger.info(\"use local dataset\")\n",
    "    data_files = {}\n",
    "    if args.train_data_dir is not None:\n",
    "        data_files[\"train\"] = os.path.join(args.train_data_dir, \"**\")\n",
    "    dataset = load_dataset(\n",
    "        \"imagefolder\",\n",
    "        data_files=data_files,\n",
    "        cache_dir=args.cache_dir,\n",
    "    )\n",
    "\n",
    "# Preprocessing the datasets.\n",
    "# We need to tokenize inputs and targets.\n",
    "column_names = dataset[\"train\"].column_names\n",
    "print(dataset[\"train\"][0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 6. Get the column names for input/target.\n",
    "dataset_columns = DATASET_NAME_MAPPING.get(args.dataset_name, None)\n",
    "if args.image_column is None:\n",
    "    image_column = dataset_columns[0] if dataset_columns is not None else column_names[0]\n",
    "else:\n",
    "    image_column = args.image_column\n",
    "    if image_column not in column_names:\n",
    "        raise ValueError(\n",
    "            f\"--image_column' value '{args.image_column}' needs to be one of: {', '.join(column_names)}\"\n",
    "        )\n",
    "if args.caption_column is None:\n",
    "    caption_column = dataset_columns[1] if dataset_columns is not None else column_names[1]\n",
    "else:\n",
    "    caption_column = args.caption_column\n",
    "    if caption_column not in column_names:\n",
    "        raise ValueError(\n",
    "            f\"--caption_column' value '{args.caption_column}' needs to be one of: {', '.join(column_names)}\"\n",
    "        )\n",
    "\n",
    "# Preprocessing the datasets.\n",
    "# We need to tokenize input captions and transform the images.\n",
    "def tokenize_captions(examples, is_train=True):\n",
    "    captions = []\n",
    "    for caption in examples[caption_column]:\n",
    "        if isinstance(caption, str):\n",
    "            captions.append(caption)\n",
    "        elif isinstance(caption, (list, np.ndarray)):\n",
    "            # take a random caption if there are multiple\n",
    "            captions.append(random.choice(caption) if is_train else caption[0])\n",
    "        else:\n",
    "            raise ValueError(\n",
    "                f\"Caption column `{caption_column}` should contain either strings or lists of strings.\"\n",
    "            )\n",
    "    inputs = tokenizer(\n",
    "        captions, max_length=tokenizer.model_max_length, padding=\"max_length\", truncation=True, return_tensors=\"pt\"\n",
    "    )\n",
    "    return inputs.input_ids\n",
    "\n",
    "# Preprocessing the datasets.\n",
    "interp_mode = transforms.InterpolationMode.BILINEAR # transforms.InterpolationMode.NEAREST\n",
    "train_transforms = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize(args.resolution, interpolation=interp_mode),\n",
    "        transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution),\n",
    "        transforms.RandomHorizontalFlip() if args.random_flip else transforms.Lambda(lambda x: x),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.5], [0.5]),\n",
    "    ]\n",
    ")\n",
    "\n",
    "def unwrap_model(model):\n",
    "    model = accelerator.unwrap_model(model)\n",
    "    model = model._orig_mod if is_compiled_module(model) else model\n",
    "    return model\n",
    "\n",
    "def preprocess_train(examples):\n",
    "    images = [image.convert(\"RGB\") for image in examples[image_column]]\n",
    "    examples[\"pixel_values\"] = [train_transforms(image) for image in images] # transform image\n",
    "    examples[\"input_ids\"] = tokenize_captions(examples) # tokenize text\n",
    "    return examples\n",
    "\n",
    "with accelerator.main_process_first():\n",
    "    if args.max_train_samples is not None:\n",
    "        dataset[\"train\"] = dataset[\"train\"].shuffle(seed=args.seed).select(range(args.max_train_samples))\n",
    "    # Set the training transforms\n",
    "    train_dataset = dataset[\"train\"].with_transform(preprocess_train)\n",
    "\n",
    "# 钩子函数，处理将数据处理成可以进行cfg的形状\n",
    "# (1,77,748)-->(2,77,748)\n",
    "def collate_fn(examples):\n",
    "    pixel_values = torch.stack([example[\"pixel_values\"] for example in examples])\n",
    "    pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()\n",
    "    input_ids = torch.stack([example[\"input_ids\"] for example in examples])\n",
    "    return {\"pixel_values\": pixel_values, \"input_ids\": input_ids}\n",
    "\n",
    "# DataLoaders creation:\n",
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "    train_dataset,\n",
    "    shuffle=True,\n",
    "    collate_fn=collate_fn,\n",
    "    batch_size=args.train_batch_size,\n",
    "    num_workers=args.dataloader_num_workers,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Scheduler and math around the number of training steps.\n",
    "num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes\n",
    "if args.max_train_steps is None:\n",
    "    # 这里计算理论上一个设备分多少数据\n",
    "    len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes)\n",
    "    # 单设备上，一个epoch中进行几次梯度更新\n",
    "    num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps)\n",
    "    # 所有设备上，总共要进行几次梯度更新\n",
    "    num_training_steps_for_scheduler = (\n",
    "        args.num_train_epochs * num_update_steps_per_epoch * accelerator.num_processes\n",
    "    )\n",
    "else:\n",
    "    num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes\n",
    "\n",
    "lr_scheduler = get_scheduler(\n",
    "    args.lr_scheduler,\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=num_warmup_steps_for_scheduler,\n",
    "    num_training_steps=num_training_steps_for_scheduler,\n",
    ")\n",
    "\n",
    "# Prepare everything with our `accelerator`.\n",
    "unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(\n",
    "    unet, optimizer, train_dataloader, lr_scheduler\n",
    ")\n",
    "\n",
    "# We need to recalculate our total training steps as the size of the training dataloader may have changed.\n",
    "# 单设备上，一个epoch中进行几次梯度更新（这里的dataloader是实际拆分后的，一个设备上实际有多少数据）\n",
    "num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)\n",
    "if args.max_train_steps is None:\n",
    "    args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch\n",
    "    if num_training_steps_for_scheduler != args.max_train_steps * accelerator.num_processes:\n",
    "        logger.warning(\n",
    "            f\"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match \"\n",
    "            f\"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. \"\n",
    "            f\"This inconsistency may result in the learning rate scheduler not functioning properly.\"\n",
    "        )\n",
    "# Afterwards we recalculate our number of training epochs\n",
    "args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)\n",
    "\n",
    "# We need to initialize the trackers we use, and also store our configuration.\n",
    "# The trackers initializes automatically on the main process.\n",
    "if accelerator.is_main_process:\n",
    "    accelerator.init_trackers(\"text2image-fine-tune\", config=vars(args))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "04/05/2025 16:42:26 - INFO - __main__ - ***** Running training *****\n",
      "04/05/2025 16:42:26 - INFO - __main__ -   Num examples = 2\n",
      "04/05/2025 16:42:26 - INFO - __main__ -   Num Epochs = 100\n",
      "04/05/2025 16:42:26 - INFO - __main__ -   Instantaneous batch size per device = 16\n",
      "04/05/2025 16:42:26 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 16\n",
      "04/05/2025 16:42:26 - INFO - __main__ -   Gradient Accumulation steps = 1\n",
      "04/05/2025 16:42:26 - INFO - __main__ -   Total optimization steps = 100\n",
      "Steps: 100%|██████████| 100/100 [18:56<00:00, 11.37s/it, lr=0.0001, step_loss=0.0136]\n",
      "Steps: 100%|██████████| 100/100 [01:21<00:00,  1.21it/s, lr=0.0001, step_loss=0.0116]"
     ]
    }
   ],
   "source": [
    "# Train!\n",
    "total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps\n",
    "\n",
    "logger.info(\"***** Running training *****\")\n",
    "logger.info(f\"  Num examples = {len(train_dataset)}\")\n",
    "logger.info(f\"  Num Epochs = {args.num_train_epochs}\")\n",
    "logger.info(f\"  Instantaneous batch size per device = {args.train_batch_size}\")\n",
    "logger.info(f\"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}\")\n",
    "logger.info(f\"  Gradient Accumulation steps = {args.gradient_accumulation_steps}\")\n",
    "logger.info(f\"  Total optimization steps = {args.max_train_steps}\")\n",
    "global_step = 0\n",
    "first_epoch = 0\n",
    "\n",
    "# Potentially load in the weights and states from a previous save\n",
    "if args.resume_from_checkpoint:\n",
    "    if args.resume_from_checkpoint != \"latest\":\n",
    "        path = os.path.basename(args.resume_from_checkpoint)\n",
    "    else:\n",
    "        # Get the most recent checkpoint\n",
    "        dirs = os.listdir(args.output_dir)\n",
    "        dirs = [d for d in dirs if d.startswith(\"checkpoint\")]\n",
    "        dirs = sorted(dirs, key=lambda x: int(x.split(\"-\")[1]))\n",
    "        path = dirs[-1] if len(dirs) > 0 else None\n",
    "\n",
    "    if path is None:\n",
    "        accelerator.print(\n",
    "            f\"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run.\"\n",
    "        )\n",
    "        args.resume_from_checkpoint = None\n",
    "        initial_global_step = 0\n",
    "    else:\n",
    "        accelerator.print(f\"Resuming from checkpoint {path}\")\n",
    "        accelerator.load_state(os.path.join(args.output_dir, path))\n",
    "        global_step = int(path.split(\"-\")[1])\n",
    "\n",
    "        initial_global_step = global_step\n",
    "        first_epoch = global_step // num_update_steps_per_epoch\n",
    "else:\n",
    "    initial_global_step = 0\n",
    "\n",
    "# 进度条\n",
    "progress_bar = tqdm(\n",
    "    range(0, args.max_train_steps),\n",
    "    initial=initial_global_step,\n",
    "    desc=\"Steps\",\n",
    "    # Only show the progress bar once on each machine.\n",
    "    disable=not accelerator.is_local_main_process,\n",
    ")\n",
    "\n",
    "for epoch in range(first_epoch, args.num_train_epochs):\n",
    "    unet.train()\n",
    "    train_loss = 0.0\n",
    "    for step, batch in enumerate(train_dataloader):\n",
    "        with accelerator.accumulate(unet):\n",
    "            # Convert images to latent space\n",
    "            latents = vae.encode(batch[\"pixel_values\"].to(dtype=weight_dtype)).latent_dist.sample()\n",
    "            latents = latents * vae.config.scaling_factor\n",
    "\n",
    "            # Sample noise that we'll add to the latents\n",
    "            noise = torch.randn_like(latents)\n",
    "            if args.noise_offset: # TODO\n",
    "                # https://www.crosslabs.org//blog/diffusion-with-offset-noise\n",
    "                noise += args.noise_offset * torch.randn(\n",
    "                    (latents.shape[0], latents.shape[1], 1, 1), device=latents.device\n",
    "                )\n",
    "\n",
    "            bsz = latents.shape[0]\n",
    "            # Sample a random timestep for each image\n",
    "            timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)\n",
    "            timesteps = timesteps.long()\n",
    "\n",
    "            # Add noise to the latents according to the noise magnitude at each timestep\n",
    "            # (this is the forward diffusion process)\n",
    "            noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)\n",
    "\n",
    "            # Get the text embedding for conditioning\n",
    "            encoder_hidden_states = text_encoder(batch[\"input_ids\"], return_dict=False)[0]\n",
    "\n",
    "            # Get the target for loss depending on the prediction type\n",
    "            if args.prediction_type is not None:\n",
    "                # set prediction_type of scheduler if defined\n",
    "                noise_scheduler.register_to_config(prediction_type=args.prediction_type)\n",
    "\n",
    "            # 预测类型：noise or v_pred\n",
    "            if noise_scheduler.config.prediction_type == \"epsilon\":\n",
    "                target = noise\n",
    "            elif noise_scheduler.config.prediction_type == \"v_prediction\":\n",
    "                target = noise_scheduler.get_velocity(latents, noise, timesteps)\n",
    "            else:\n",
    "                raise ValueError(f\"Unknown prediction type {noise_scheduler.config.prediction_type}\")\n",
    "\n",
    "            # Predict the noise residual and compute loss\n",
    "            model_pred = unet(noisy_latents, timesteps, encoder_hidden_states, return_dict=False)[0]\n",
    "            \n",
    "            # 是否采用SNR-gamma策略加快模型收敛\n",
    "            if args.snr_gamma is None:\n",
    "                loss = F.mse_loss(model_pred.float(), target.float(), reduction=\"mean\")\n",
    "            else:\n",
    "                # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.\n",
    "                # Since we predict the noise instead of x_0, the original formulation is slightly changed.\n",
    "                # This is discussed in Section 4.2 of the same paper.\n",
    "                snr = compute_snr(noise_scheduler, timesteps)\n",
    "                mse_loss_weights = torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(\n",
    "                    dim=1\n",
    "                )[0]\n",
    "                if noise_scheduler.config.prediction_type == \"epsilon\":\n",
    "                    mse_loss_weights = mse_loss_weights / snr\n",
    "                elif noise_scheduler.config.prediction_type == \"v_prediction\":\n",
    "                    mse_loss_weights = mse_loss_weights / (snr + 1)\n",
    "\n",
    "                loss = F.mse_loss(model_pred.float(), target.float(), reduction=\"none\")\n",
    "                loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights\n",
    "                loss = loss.mean()\n",
    "\n",
    "            # Gather the losses across all processes for logging (if we use distributed training).\n",
    "            avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()\n",
    "            train_loss += avg_loss.item() / args.gradient_accumulation_steps\n",
    "\n",
    "            # Backpropagate\n",
    "            accelerator.backward(loss)\n",
    "            if accelerator.sync_gradients:\n",
    "                params_to_clip = lora_layers\n",
    "                accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "        # Checks if the accelerator has performed an optimization step behind the scenes\n",
    "        if accelerator.sync_gradients:\n",
    "            progress_bar.update(1)\n",
    "            global_step += 1\n",
    "            accelerator.log({\"train_loss\": train_loss}, step=global_step)\n",
    "            train_loss = 0.0\n",
    "            # 达到指定step,保存模型\n",
    "            if global_step % args.checkpointing_steps == 0:\n",
    "                if accelerator.is_main_process:\n",
    "                    # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`\n",
    "                    if args.checkpoints_total_limit is not None:\n",
    "                        checkpoints = os.listdir(args.output_dir)\n",
    "                        checkpoints = [d for d in checkpoints if d.startswith(\"checkpoint\")]\n",
    "                        checkpoints = sorted(checkpoints, key=lambda x: int(x.split(\"-\")[1]))\n",
    "\n",
    "                        # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints\n",
    "                        if len(checkpoints) >= args.checkpoints_total_limit:\n",
    "                            num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1\n",
    "                            removing_checkpoints = checkpoints[0:num_to_remove]\n",
    "\n",
    "                            logger.info(\n",
    "                                f\"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints\"\n",
    "                            )\n",
    "                            logger.info(f\"removing checkpoints: {', '.join(removing_checkpoints)}\")\n",
    "\n",
    "                            for removing_checkpoint in removing_checkpoints:\n",
    "                                removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)\n",
    "                                shutil.rmtree(removing_checkpoint)\n",
    "\n",
    "                    save_path = os.path.join(args.output_dir, f\"checkpoint-{global_step}\")\n",
    "                    accelerator.save_state(save_path)\n",
    "\n",
    "                    unwrapped_unet = unwrap_model(unet)\n",
    "                    unet_lora_state_dict = convert_state_dict_to_diffusers(\n",
    "                        get_peft_model_state_dict(unwrapped_unet)\n",
    "                    )\n",
    "\n",
    "                    StableDiffusionPipeline.save_lora_weights(\n",
    "                        save_directory=save_path,\n",
    "                        unet_lora_layers=unet_lora_state_dict,\n",
    "                        safe_serialization=True,\n",
    "                    )\n",
    "\n",
    "                    logger.info(f\"Saved state to {save_path}\")\n",
    "\n",
    "        logs = {\"step_loss\": loss.detach().item(), \"lr\": lr_scheduler.get_last_lr()[0]}\n",
    "        progress_bar.set_postfix(**logs)\n",
    "\n",
    "        if global_step >= args.max_train_steps:\n",
    "            break\n",
    "\n",
    "    # 一个epoch结束，验证效果\n",
    "    if accelerator.is_main_process:\n",
    "        if args.validation_prompt is not None and epoch % args.validation_epochs == 0:\n",
    "            # create pipeline\n",
    "            pipeline = DiffusionPipeline.from_pretrained(\n",
    "                args.pretrained_model_name_or_path,\n",
    "                unet=unwrap_model(unet),\n",
    "                revision=args.revision,\n",
    "                variant=args.variant,\n",
    "                torch_dtype=weight_dtype,\n",
    "            )\n",
    "            images = log_validation(pipeline, args, accelerator, epoch)\n",
    "\n",
    "            del pipeline\n",
    "            torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2025-04-05 16:44:09,489] [INFO] [real_accelerator.py:161:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in sd-model-finetuned-lora/pytorch_lora_weights.safetensors\n"
     ]
    }
   ],
   "source": [
    "# Save the lora layers\n",
    "accelerator.wait_for_everyone()\n",
    "if accelerator.is_main_process:\n",
    "    unet = unet.to(torch.float32)\n",
    "\n",
    "    unwrapped_unet = unwrap_model(unet)\n",
    "    unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unwrapped_unet))\n",
    "    StableDiffusionPipeline.save_lora_weights(\n",
    "        save_directory=args.output_dir,\n",
    "        unet_lora_layers=unet_lora_state_dict,\n",
    "        safe_serialization=True,\n",
    "    )\n",
    "\n",
    "    # Final inference\n",
    "    # Load previous pipeline\n",
    "    if args.validation_prompt is not None:\n",
    "        pipeline = DiffusionPipeline.from_pretrained(\n",
    "            args.pretrained_model_name_or_path,\n",
    "            revision=args.revision,\n",
    "            variant=args.variant,\n",
    "            torch_dtype=weight_dtype,\n",
    "        )\n",
    "\n",
    "        # load attention processors\n",
    "        pipeline.load_lora_weights(args.output_dir)\n",
    "\n",
    "        # run inference\n",
    "        images = log_validation(pipeline, args, accelerator, epoch, is_final_validation=True)\n",
    "\n",
    "    # 上传模型到Hub\n",
    "    if args.push_to_hub:\n",
    "        save_model_card(\n",
    "            repo_id,\n",
    "            images=images,\n",
    "            base_model=args.pretrained_model_name_or_path,\n",
    "            dataset_name=args.dataset_name,\n",
    "            repo_folder=args.output_dir,\n",
    "        )\n",
    "        upload_folder(\n",
    "            repo_id=repo_id,\n",
    "            folder_path=args.output_dir,\n",
    "            commit_message=\"End of training\",\n",
    "            ignore_patterns=[\"step_*\", \"epoch_*\"],\n",
    "        )\n",
    "\n",
    "accelerator.end_training()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
